{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Continuous Training with Kubeflow Pipeline and Vertex AI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Learning Objectives:**\n",
    "1. Learn how to use KF pre-built components\n",
    "1. Learn how to use KF lightweight python components\n",
    "1. Learn how to build a KF pipeline with these components\n",
    "1. Learn how to compile, upload, and run a KF pipeline\n",
    "\n",
    "\n",
    "In this lab, you will build, deploy, and run a KFP pipeline that orchestrates the **Vertex AI** services to train, tune, and deploy a **scikit-learn** model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import aiplatform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "REGION = 'us-central1'\n",
    "PROJECT_ID = !(gcloud config get-value project)\n",
    "PROJECT_ID = PROJECT_ID[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set `PATH` to include the directory containing KFP CLI\n",
    "PATH=%env PATH\n",
    "%env PATH=/home/jupyter/.local/bin:{PATH}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding the pipeline design\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The workflow implemented by the pipeline is defined using a Python based Domain Specific Language (DSL). The pipeline's DSL is in the `pipeline_vertex/pipeline.py` file that we will generate below.\n",
    "\n",
    "The pipeline's DSL has been designed to avoid hardcoding any environment specific settings like file paths or connection strings. These settings are provided to the pipeline code through a set of environment variables.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build the trainer image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The training step in the pipeline will require a custom training container. The custom training image is defined in `trainer_image/Dockerfile`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat trainer_image_vertex/Dockerfile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now build and push this trainer container to the container registry:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_NAME='trainer_image_covertype_vertex'\n",
    "TAG='latest'\n",
    "TRAINING_CONTAINER_IMAGE_URI=f'gcr.io/{PROJECT_ID}/{IMAGE_NAME}:{TAG}'\n",
    "TRAINING_CONTAINER_IMAGE_URI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud builds submit --timeout 15m --tag $TRAINING_CONTAINER_IMAGE_URI trainer_image_vertex"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To match the ml framework version we use at training time while serving the model, we will have to supply the following serving container to the pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SERVING_CONTAINER_IMAGE_URI = 'us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-20:latest'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** If you change the version of the training ml framework you'll have to supply a serving container with matchin version (see [pre-built containers for prediction](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers))."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building and deploying the pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us write the pipeline to disk:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./pipeline_vertex/pipeline.py\n",
    "# Copyright 2021 Google LLC\n",
    "\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this\n",
    "# file except in compliance with the License. You may obtain a copy of the License at\n",
    "\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "    \n",
    "# Unless required by applicable law or agreed to in writing, software \n",
    "# distributed under the License is distributed on an \"AS IS\"\n",
    "# BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either \n",
    "# express or implied. See the License for the specific language governing \n",
    "# permissions and limitations under the License.\n",
    "\"\"\"Kubeflow Covertype Pipeline.\"\"\"\n",
    "import os\n",
    "\n",
    "from kfp import dsl\n",
    "from kfp.components import create_component_from_func_v2\n",
    "\n",
    "from tuning_lightweight_component import tune_hyperparameters\n",
    "from training_lightweight_component import train_and_deploy\n",
    "\n",
    "\n",
    "PIPELINE_ROOT = os.getenv('PIPELINE_ROOT')\n",
    "PROJECT_ID = os.getenv('PROJECT_ID')\n",
    "REGION = os.getenv('REGION')\n",
    "\n",
    "TRAINING_CONTAINER_IMAGE_URI = os.getenv('TRAINING_CONTAINER_IMAGE_URI')\n",
    "SERVING_CONTAINER_IMAGE_URI = os.getenv('SERVING_CONTAINER_IMAGE_URI')\n",
    "\n",
    "TRAINING_FILE_PATH = os.getenv('TRAINING_FILE_PATH')\n",
    "VALIDATION_FILE_PATH = os.getenv('VALIDATION_FILE_PATH')\n",
    "\n",
    "MAX_TRIAL_COUNT = os.getenv('MAX_TRIAL_COUNT', 5)\n",
    "PARALLEL_TRIAL_COUNT = os.getenv('PARALLEL_TRIAL_COUNT', 5)\n",
    "THRESHOLD = os.getenv('THRESHOLD', 0.6)\n",
    "\n",
    "\n",
    "tune_hyperparameters_component = create_component_from_func_v2(\n",
    "    tune_hyperparameters,\n",
    "    base_image='python:3.8',\n",
    "    output_component_file='covertype_kfp_tune_hyperparameters.yaml',\n",
    "    packages_to_install=['google-cloud-aiplatform'],\n",
    ")\n",
    "\n",
    "\n",
    "train_and_deploy_component = create_component_from_func_v2(\n",
    "    train_and_deploy,\n",
    "    base_image='python:3.8',\n",
    "    output_component_file='covertype_kfp_train_and_deploy.yaml',\n",
    "    packages_to_install=['google-cloud-aiplatform'],\n",
    "\n",
    ")\n",
    "\n",
    "\n",
    "@dsl.pipeline(\n",
    "    name=\"covertype-kfp-pipeline\",\n",
    "    description=\"The pipeline training and deploying the Covertype classifier\",\n",
    "    pipeline_root=PIPELINE_ROOT,\n",
    ")\n",
    "def covertype_train(\n",
    "    training_container_uri: str = TRAINING_CONTAINER_IMAGE_URI,\n",
    "    serving_container_uri: str = SERVING_CONTAINER_IMAGE_URI,\n",
    "    training_file_path: str = TRAINING_FILE_PATH,\n",
    "    validation_file_path: str = VALIDATION_FILE_PATH,\n",
    "    accuracy_deployment_threshold: float = THRESHOLD,\n",
    "    max_trial_count: int = MAX_TRIAL_COUNT,\n",
    "    parallel_trial_count: int = PARALLEL_TRIAL_COUNT,\n",
    "    pipeline_root: str = PIPELINE_ROOT,\n",
    "):\n",
    "    staging_bucket = f'{pipeline_root}/staging'\n",
    "    \n",
    "    tuning_op = tune_hyperparameters_component(\n",
    "        project=PROJECT_ID,\n",
    "        location=REGION,\n",
    "        container_uri=training_container_uri,\n",
    "        training_file_path=training_file_path,\n",
    "        validation_file_path=validation_file_path,\n",
    "        staging_bucket=staging_bucket,\n",
    "        max_trial_count=max_trial_count,\n",
    "        parallel_trial_count=parallel_trial_count,\n",
    "    )\n",
    "    \n",
    "    accuracy = tuning_op.outputs['best_accuracy']\n",
    "    \n",
    "    with dsl.Condition(accuracy >= accuracy_deployment_threshold, name=\"deploy_decision\"):    \n",
    "        train_and_deploy_op = train_and_deploy_component(\n",
    "            project=PROJECT_ID,\n",
    "            location=REGION,\n",
    "            container_uri=training_container_uri,\n",
    "            serving_container_uri=serving_container_uri,\n",
    "            training_file_path=training_file_path,\n",
    "            validation_file_path=validation_file_path,\n",
    "            staging_bucket=staging_bucket,\n",
    "            alpha=tuning_op.outputs['best_alpha'], \n",
    "            max_iter=tuning_op.outputs['best_max_iter'],            \n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compile the pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let stat by defining the environment variables that will be passed to the pipeline compiler:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ARTIFACT_STORE = f'gs://{PROJECT_ID}-vertex'\n",
    "PIPELINE_ROOT = f'{ARTIFACT_STORE}/pipeline'\n",
    "DATA_ROOT = f'{ARTIFACT_STORE}/data'\n",
    "\n",
    "TRAINING_FILE_PATH = f'{DATA_ROOT}/training/dataset.csv'\n",
    "VALIDATION_FILE_PATH = f'{DATA_ROOT}/validation/dataset.csv'\n",
    "\n",
    "%env PIPELINE_ROOT={PIPELINE_ROOT}\n",
    "%env PROJECT_ID={PROJECT_ID}\n",
    "%env REGION={REGION}\n",
    "%env SERVING_CONTAINER_IMAGE_URI={SERVING_CONTAINER_IMAGE_URI}\n",
    "%env TRAINING_CONTAINER_IMAGE_URI={TRAINING_CONTAINER_IMAGE_URI}\n",
    "%env TRAINING_FILE_PATH={TRAINING_FILE_PATH}\n",
    "%env VALIDATION_FILE_PATH={VALIDATION_FILE_PATH}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us make sure that the `ARTIFACT_STORE` has been created, and let us create it if not:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gsutil ls | grep ^{ARTIFACT_STORE}/$ || gsutil mb -l {REGION} {ARTIFACT_STORE}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** In case the artifact store was not created and properly set before hand, you may need\n",
    "to run in **CloudShell** the following command to allow Vertex AI to access it:\n",
    "\n",
    "```\n",
    "PROJECT_ID=$(gcloud config get-value project)\n",
    "PROJECT_NUMBER=$(gcloud projects list --filter=\"name=$PROJECT_ID\" --format=\"value(PROJECT_NUMBER)\")\n",
    "gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
    "    --member=\"serviceAccount:$PROJECT_NUMBER-compute@developer.gserviceaccount.com\" \\\n",
    "    --role=\"roles/storage.objectAdmin\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Use the CLI compiler to compile the pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We compile the pipeline from the Python file we generated into a JSON description using the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PIPELINE_JSON = 'covertype_kfp_pipeline.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!dsl-compile-v2 --py pipeline_vertex/pipeline.py --output $PIPELINE_JSON"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** You can also use the Python SDK to compile the pipeline from its python function\n",
    "\n",
    "```python\n",
    "compiler.Compiler().compile(\n",
    "    pipeline_func=covertype_train, \n",
    "    package_path=PIPELINE_JSON,\n",
    ")\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is the pipeline file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head {PIPELINE_JSON}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deploy the pipeline package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
    "\n",
    "pipeline = aiplatform.PipelineJob(\n",
    "    display_name='covertype_kfp_pipeline', \n",
    "    template_path=PIPELINE_JSON, \n",
    "    enable_caching=False,\n",
    ")\n",
    "\n",
    "pipeline.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2021 Google LLC\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "    https://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License."
   ]
  }
 ],
 "metadata": {
  "environment": {
   "name": "tf2-gpu.2-3.m75",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-3:m75"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
